# Requires transformers>=4.51.0
# Requires sentence-transformers>=2.7.0

import os
from sentence_transformers import SentenceTransformer, CrossEncoder
import config

def set_environment():
    """设置环境变量"""
    os.environ["HF_ENDPOINT"] = config.HF_ENDPOINT
    os.environ["HF_HUB_DISABLE_XET"] = config.HF_HUB_DISABLE_XET
    
def load_embedding_model(model_config=None):
    """
    加载Qwen3 Embedding模型
    
    Args:
        model_config: 模型配置，如果为None则使用config.py中的默认配置
        
    Returns:
        加载的模型
    """
    if model_config is None:
        model_config = config.EMBEDDING_CONFIG
    
    # 根据配置确定模型路径
    if model_config["local_model_path"]:
        print(f"正在加载本地模型: {model_config['local_model_path']}")
        model_path = model_config["local_model_path"]
    else:
        print(f"正在从HuggingFace加载模型: {model_config['model_id']}")
        model_path = model_config["model_id"]
    
    # 如果使用CrossEncoder加载方式
    if model_config.get("use_cross_encoder", False):
        print("使用CrossEncoder加载模型")
        model = CrossEncoder(model_path)
        return model
    
    # 根据是否使用flash_attention_2加载模型
    if model_config["use_flash_attention"]:
        model = SentenceTransformer(
            model_path,
            model_kwargs={"attn_implementation": "flash_attention_2", "device_map": "auto"},
            tokenizer_kwargs={"padding_side": "left"},
        )
        print("已启用flash_attention_2加速")
    else:
        model = SentenceTransformer(model_path)
    
    return model

def generate_embeddings(model, texts):
    """
    生成文本嵌入向量
    
    Args:
        model: 加载的模型
        texts: 文本列表
        
    Returns:
        文本嵌入向量
    """
    return model.encode(texts)

def calculate_similarity(model, embeddings1, embeddings2):
    """
    计算两组嵌入向量之间的相似度
    
    Args:
        model: 加载的模型
        embeddings1: 第一组嵌入向量
        embeddings2: 第二组嵌入向量
        
    Returns:
        相似度矩阵
    """
    return model.similarity(embeddings1, embeddings2)

def run_example():
    """运行示例"""
    # 设置环境变量
    set_environment()
    
    # 加载模型
    model = load_embedding_model()
    
    # 生成嵌入向量
    query_embeddings = generate_embeddings(model, config.EXAMPLE_QUERIES)
    document_embeddings = generate_embeddings(model, config.EXAMPLE_DOCUMENTS)
    
    # 计算相似度
    similarity = calculate_similarity(model, query_embeddings, document_embeddings)
    
    # 打印结果
    print("相似度矩阵:")
    print(similarity)

if __name__ == "__main__":
    run_example()
